import subprocess
import rrd_base
import re
import time

import color

try:
    from pyrrd.rrd import DataSource, RRA, RRD
    from pyrrd.graph import DEF, CDEF, VDEF
    from pyrrd.graph import LINE, AREA, GPRINT
    from pyrrd.graph import ColorAttributes, Graph
except:
    print "Couldn't load required modules. You probably need to install it with:"
    sys.exit('sudo apt-get install python-pyrrd python-pysnmp4')

class RrdCpuUsage(rrd_base.RrdBase):
    def __init__(self, domain, top_lines):
        super(RrdCpuUsage,self).__init__(domain=domain, group="cpu", instance="cpuusage")
        self.top_lines = top_lines

    def create_file(self):
        dataSources = []
        for ds in ['us', 'sy', 'ni', 'id', 'wa', 'hi', 'si', 'st']:
            dataSources.append(DataSource( dsName=ds, dsType='GAUGE', heartbeat=self.heartbeat()))

        roundRobinArchives = []
        for func in ['AVERAGE', 'MAX', 'MIN']:
            for step,row in self.rra_pattern():
                roundRobinArchives.append(RRA(cf=func, xff=0.5, steps=step, rows=row))

        self.myRRD = RRD(self.filename, ds=dataSources, rra=roundRobinArchives, step=self.data_step())
        self.myRRD.create()

# %Cpu(s):  6.4 us,  3.2 sy,  0.0 ni, 88.7 id,  0.6 wa,  0.0 hi,  1.1 si,  0.0 st

    def update(self):
        sums = [0, 0, 0, 0, 0, 0, 0, 0]
        samples = 0
        for line in self.top_lines:
            # Must use search, some architectures prefix CPU with a %
            if re.search("Cpu\(s\): ", line):
                samples += 1
                i = 0
                for sample in re.sub("[^0-9\. ]", "", line).split():
                    sums[i] += float(sample)
                    i += 1

        for sum in sums:
            self.myRRD.bufferValue(int(round(time.time())),
                                   sums[0]/samples,
                                   sums[1]/samples,
                                   sums[2]/samples,
                                   sums[3]/samples,
                                   sums[4]/samples,
                                   sums[5]/samples,
                                   sums[6]/samples,
                                   sums[7]/samples)
            self.myRRD.update()
            break

    def graph(self):
        def_us    = DEF(rrdfile=self.myRRD.filename, vname='us',   cdef="AVERAGE", dsName='us')
        def_sy    = DEF(rrdfile=self.myRRD.filename, vname='sy',   cdef="AVERAGE", dsName='sy')
        def_ni    = DEF(rrdfile=self.myRRD.filename, vname='ni',   cdef="AVERAGE", dsName='ni')
        def_id    = DEF(rrdfile=self.myRRD.filename, vname='id',   cdef="AVERAGE", dsName='id')
        def_wa    = DEF(rrdfile=self.myRRD.filename, vname='wa',   cdef="AVERAGE", dsName='wa')
        def_hi    = DEF(rrdfile=self.myRRD.filename, vname='hi',   cdef="AVERAGE", dsName='hi')
        def_si    = DEF(rrdfile=self.myRRD.filename, vname='si',   cdef="AVERAGE", dsName='si')
        def_st    = DEF(rrdfile=self.myRRD.filename, vname='st',   cdef="AVERAGE", dsName='st')

        area_us   = AREA(defObj=def_us,  color=color.RGBA['navy'], legend="us", stack=True)
        area_sy   = AREA(defObj=def_sy,  color=color.RGBA['red'],   legend="sy", stack=True)
        area_ni   = AREA(defObj=def_ni,  color=color.RGBA['blue'], legend="ni", stack=True)
        area_id   = AREA(defObj=def_id,  color=color.RGBA['green'], legend="id", stack=True)
        area_wa   = AREA(defObj=def_wa,  color=color.RGBA['darkorange'], legend="wa", stack=True)
        area_hi   = AREA(defObj=def_hi,  color=color.RGBA['skyblue'], legend="hi", stack=True)
        area_si   = AREA(defObj=def_si,  color=color.RGBA['purple'], legend="si", stack=True)
        area_st   = AREA(defObj=def_st,  color=color.RGBA['violet'], legend="st", stack=True)

        for tf in self.graph_time_frames().keys():
            g = Graph(self.image_name(tf),
                      start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                      rigid=True,
                      title='"CPU Usage on {0} - Updated {1}"'.format(self.data_domain, time.strftime("%Y%m%d at %H:%M")),
                      height=self.graph_height, width=self.graph_width,
                      lower_limit=0.001, upper_limit=105,
                      x_grid=self.grid[tf],
                      step=self.graph_step_times[tf])
            g.data.extend([  def_us,  def_sy,  def_ni,  def_id,  def_wa,  def_hi,  def_si,  def_st,
                            area_us, area_sy, area_ni, area_id, area_wa, area_hi, area_si, area_st
                          ])
            g.write()

            if False:
                g = Graph(self.icon_name(tf),
                          start=int(time.time())-self.graph_time_frames()[tf], end=int(time.time()),
                          lower_limit=0.001, upper_limit=105,
                          rigid=True, no_legend=True,
                          height=self.icon_height, width=self.icon_width,
                          step=self.graph_step_times[tf])
                g.data.extend([  def_us,  def_sy,  def_ni,  def_id,  def_wa,  def_hi,  def_si,  def_st,
                                area_us, area_sy, area_ni, area_id, area_wa, area_hi, area_si, area_st
                              ])
                g.write()
